from cion.core.hardware.rtmq.RTMQ_StdLib import *

C_CION_SEQR = StdCore(16, 0, ["LED", "TTL0", "TTL1", "ECTR"] +
                      [f"C{i:02}" for i in range(32)])

# Change "COM3" to the corresponding COM port
seqr = UARTDevice(["COM7"], 4000000)

# -------- Port Assignment --------

"""
External trigger input: J2 (SMA)

TTL channel 0:
    port  0 -  7: J4  (pin with dot mark is port  0)
    port  8 - 15: J10 (pin with dot mark is port  8)
    port 16 - 23: J5  (pin with dot mark is port 16)
    port 24 - 31: J11 (pin with dot mark is port 24)

TTL channel 1:
    port  0 -  7: J6  (pin with dot mark is port  0)
    port  8 - 15: J8  (pin with dot mark is port  8)
    port 16 - 23: J7  (pin with dot mark is port 16)
    port 24 - 31: J9  (pin with dot mark is port 24)

Counter channel:
     0 -  7: J12 (pin with dot mark is port  0)
     8 - 15: J13 (pin with dot mark is port  8)
    16 - 23: J14 (pin with dot mark is port 16)
    24 - 31: J15 (pin with dot mark is port 24)

"""

# -------- Peripheral API --------

def TTL_Stage(chn, ttl, dur):
    """
    
    Append a TTL output stage to current flow.

    Parameters
    ----------
    chn : int
        TTL output channel, can be 0 or 1, each contains 32 ports.
        chn == 2 is reserved for LED output.
    ttl : int
        Output state of this stage, each bit corresponds to a port.
    dur : int OR float
        Duration of this stage. <dur> is in micro-sec if it is
        a float number, otherwise <dur> is in clock cycles (5ns).

    Returns
    -------
    None.

    """

    if isinstance(dur, float):
        dur = int(dur * 200)
    StartTimedTask(dur)
    if chn == 0:
        Set(R.TTL0, ttl)
    elif chn == 1:
        Set(R.TTL1, ttl)
    else:
        Set(R.LED, ttl)

def ExtTrigger():
    """
    
    Wait for external trigger.

    Returns
    -------
    None.

    """
    WaitTrigger(2)

def Cntr_ON(cntr):
    """
    
    Reset and enable gated counter channels.

    Parameters
    ----------
    cntr : list
        List of counter channel numbers to be enabled.
        Start from 0.

    Returns
    -------
    None.

    """
    msk = sum([1 << i for i in cntr])
    StartTimedTask(20)
    ASM(f"LDH - -  TR0  {msk}")
    ASM(f"LDL - -  TR0  {msk}")
    Pass()
    ASM(f"SMK - -  ECTR !NUL  TR0")

def Cntr_OFF(cntr):
    """
    
    Disable gated counter channels and update results.

    Parameters
    ----------
    cntr : list
        List of counter channel numbers to be disabled.
        Start from 0.

    Returns
    -------
    None.

    """
    msk = sum([1 << i for i in cntr])
    StartTimedTask(20)
    ASM(f"LDH - -  TR0  {msk}")
    ASM(f"LDL - -  TR0  {msk}")
    Pass()
    ASM(f"SMK - -  ECTR  NUL  TR0")
    Pass(6)

def GetCounts(cntr):
    """
    
    Acquire counting results.

    Parameters
    ----------
    cntr : list
        List of counter channel numbers to be acquired.
        Start from 0.

    Returns
    -------
    None.

    """
    for c in cntr:
        SerialSend(eval(f"R.C{c:02}"))

def NewExperiment():
    """
    
    Declare a new experiment flow.

    Returns
    -------
    <ExpFlow>
        Instance of the experiment flow object.

    """
    return ExpFlow(C_CION_SEQR)

def Run(flow, ret=0):
    """
    
    Run an experiment flow.

    Parameters
    ----------
    flow : <ExpFlow>
        The experiment flow to be run.
    ret : int, optional
        Expected number of returned data points.
        The default is 0.

    Returns
    -------
    res : list
        List of returned data.

    """
    res = seqr.run(0, 0, flow, ret)
    if ret != 0:
        return res

# -------- Debug Helper --------

def EF(func, ret=0):
    def _wrapped(*args, **kwargs):
        f = ExpFlow(C_CION_SEQR)
        with f:
            func(*args, **kwargs)
        res = seqr.run(0, 0, f, ret)
        if ret != 0:
            return res
    return _wrapped

def CF(func, ret=0):
    def _wrapped(*args, **kwargs):
        f = CfgFlow(C_CION_SEQR)
        with f:
            func(*args, **kwargs)
        res = seqr.run(0, 0, f, ret)
        if ret != 0:
            return res
    return _wrapped

def GetReg(reg, fmt=None):
    res = CF(SerialSend, 1)(reg, False)[0]
    if fmt is not None:
        res = fmt(res)
    print(res)

def SetReg(reg, val):
    CF(Set)(reg, val)
